import streamlit as st
import pandas as pd
import qianfan
import json
import re
import random
import ast

# ==========================================
# 👇【API Key 配置区域】👇
# ==========================================
access_key = "".strip()
secret_key = "".strip()

chat_comp = qianfan.ChatCompletion(access_key=access_key, secret_key=secret_key)

# --- 1. 页面配置 ---
st.set_page_config(
    page_title="文心·刑辩智脑", 
    layout="wide", 
    page_icon="⚖️", 
    initial_sidebar_state="expanded"
)

# ==========================================
# 🎨 前端美化核心区域 (CSS Injection)
# ==========================================
st.markdown("""
<style>
    /* 全局字体与背景 */
    @import url('https://fonts.googleapis.com/css2?family=Noto+Serif+SC:wght@400;700&family=Roboto:wght@300;400;700&display=swap');
    
    .stApp {
        background-color: #f8f9fc; /* 极淡的蓝灰色背景 */
        font-family: 'Roboto', 'Noto Serif SC', sans-serif;
    }

    /* 标题样式 */
    h1, h2, h3 {
        color: #1e293b; /* 深蓝灰色 */
        font-weight: 700;
        letter-spacing: -0.5px;
    }
    
    /* 侧边栏美化 */
    [data-testid="stSidebar"] {
        background-color: #ffffff;
        border-right: 1px solid #e2e8f0;
        box-shadow: 2px 0 10px rgba(0,0,0,0.02);
    }

    /* 按钮美化 - 渐变与悬停 */
    div.stButton > button {
        width: 100%;
        border-radius: 8px;
        height: 50px;
        font-weight: 600;
        font-size: 16px;
        border: none;
        transition: all 0.3s ease;
        box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1), 0 2px 4px -1px rgba(0, 0, 0, 0.06);
    }
    
    /* 默认按钮 (Secondary) */
    div.stButton > button[kind="secondary"] {
        background: white;
        color: #1e293b;
        border: 1px solid #cbd5e1;
    }
    div.stButton > button[kind="secondary"]:hover {
        background: #f1f5f9;
        transform: translateY(-2px);
    }

    /* 主要按钮 (Primary) */
    div.stButton > button[kind="primary"] {
        background: linear-gradient(135deg, #1e293b 0%, #334155 100%); /* 深蓝渐变 */
        color: white;
    }
    div.stButton > button[kind="primary"]:hover {
        background: linear-gradient(135deg, #0f172a 0%, #1e293b 100%);
        box-shadow: 0 10px 15px -3px rgba(0, 0, 0, 0.1), 0 4px 6px -2px rgba(0, 0, 0, 0.05);
        transform: translateY(-2px);
    }

    /* 输入框美化 */
    .stTextArea textarea {
        border-radius: 10px;
        border: 1px solid #e2e8f0;
        background-color: #ffffff;
        padding: 15px;
        font-size: 15px;
        line-height: 1.6;
        box-shadow: inset 0 2px 4px 0 rgba(0, 0, 0, 0.05);
        transition: border-color 0.2s;
    }
    .stTextArea textarea:focus {
        border-color: #b49162; /* 金色聚焦 */
        box-shadow: 0 0 0 2px rgba(180, 145, 98, 0.2);
    }

    /* Tab 标签页美化 */
    .stTabs [data-baseweb="tab-list"] {
        gap: 8px;
        background-color: transparent;
    }
    .stTabs [data-baseweb="tab"] {
        height: 50px;
        white-space: pre-wrap;
        background-color: #ffffff;
        border-radius: 8px 8px 0 0;
        gap: 1px;
        padding-top: 10px;
        padding-bottom: 10px;
        border: 1px solid #e2e8f0;
        border-bottom: none;
        font-weight: 600;
    }
    .stTabs [aria-selected="true"] {
        background-color: #ffffff;
        color: #b49162 !important; /* 选中变金 */
        border-top: 3px solid #b49162 !important;
    }

    /* 聊天气泡美化 */
    .stChatMessage {
        background-color: #ffffff;
        border: 1px solid #f1f5f9;
        box-shadow: 0 2px 5px rgba(0,0,0,0.03);
        border-radius: 12px;
    }

    /* --- 自定义组件样式 --- */
    
    /* 时间轴 */
    .timeline-container {
        position: relative;
        padding: 20px 0;
    }
    .timeline-item {
        padding: 15px 20px;
        border-left: 3px solid #cbd5e1;
        margin-left: 10px;
        position: relative;
        transition: all 0.3s ease;
    }
    .timeline-item:hover {
        border-left-color: #b49162; /* 悬停变金 */
        background-color: #fffcf5;
    }
    .timeline-item::before {
        content: '';
        position: absolute;
        left: -9px; /* 线宽3px/2 + 点宽15px/2 */
        top: 20px;
        width: 15px;
        height: 15px;
        border-radius: 50%;
        background-color: #ffffff;
        border: 3px solid #b49162;
    }
    .timeline-time {
        font-family: 'Roboto Mono', monospace;
        color: #64748b;
        font-size: 0.85em;
        margin-bottom: 5px;
        font-weight: 700;
        text-transform: uppercase;
    }
    .timeline-content {
        color: #334155;
        font-size: 1.05em;
        line-height: 1.5;
        font-weight: 500;
    }

    /* 结果卡片 */
    .result-card {
        background: white;
        padding: 20px;
        border-radius: 12px;
        border: 1px solid #e2e8f0;
        border-left: 5px solid #1e293b; /* 深蓝左标 */
        box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.05);
        height: 100%;
        transition: transform 0.2s;
    }
    .result-card:hover {
        transform: translateY(-3px);
        box-shadow: 0 10px 15px -3px rgba(0, 0, 0, 0.1);
    }
    .card-title {
        color: #94a3b8;
        font-size: 0.85rem;
        text-transform: uppercase;
        letter-spacing: 1px;
        font-weight: 700;
        margin-bottom: 10px;
    }
    .card-value {
        color: #1e293b;
        font-size: 1.3rem;
        font-weight: 800;
        line-height: 1.4;
    }
    .card-value-highlight {
        color: #d9534f; /* 警告红 */
    }
</style>
""", unsafe_allow_html=True)

# --- 2. 数据加载 ---
@st.cache_data
def load_data():
    try:
        return pd.read_pickle('law_data_clean.pkl')
    except:
        return pd.DataFrame()

df = load_data()

# 加载素材库
@st.cache_data
def load_learning_materials():
    try:
        with open('素材.txt', 'r', encoding='utf-8') as f:
            content = f.read()
        return content
    except:
        return ""

materials_content = load_learning_materials()

# --- 3. 状态管理 ---
if 'current_fact' not in st.session_state:
    st.session_state.current_fact = ""
if 'analysis_result' not in st.session_state:
    st.session_state.analysis_result = None 
if 'debate_history' not in st.session_state:
    st.session_state.debate_history = []
if 'messages' not in st.session_state:
    st.session_state.messages = []

def on_case_selected():
    idx = st.session_state.selected_case_index
    if idx is not None:
        st.session_state.current_fact = df.loc[idx]['fact']
        st.session_state.analysis_result = None
        st.session_state.debate_history = [] 
        st.session_state.messages = [] 

def on_clear_click():
    st.session_state.current_fact = ""
    st.session_state.analysis_result = None
    st.session_state.debate_history = []
    st.session_state.messages = []

# --- 4. 侧边栏 ---
with st.sidebar:
    st.markdown("## ⚖️ 刑辩智囊 Pro")
    st.caption("AI 驱动的刑辩全流程辅助系统")
    st.markdown("---")
    
    if not df.empty:
        display_df = df.sample(min(100, len(df)))
        st.selectbox(
            "📚 选择模拟案例", 
            display_df.index, 
            index=None, 
            key="selected_case_index", 
            on_change=on_case_selected, 
            placeholder="点击浏览案例库..."
        )
    else:
        st.error("❌ 数据库未连接")
        
    st.markdown("---")
    
    # 侧边栏状态指示
    col_a, col_b = st.columns(2)
    with col_a:
        if materials_content:
            st.markdown("""
            <div style="text-align:center; padding:10px; background:#f0fdf4; border-radius:8px; border:1px solid #bbf7d0;">
                <div style="font-size:20px;">📚</div>
                <div style="font-size:12px; color:#166534; font-weight:bold;">素材已加载</div>
            </div>
            """, unsafe_allow_html=True)
    with col_b:
        st.markdown("""
        <div style="text-align:center; padding:10px; background:#eff6ff; border-radius:8px; border:1px solid #bfdbfe;">
            <div style="font-size:20px;">🤖</div>
            <div style="font-size:12px; color:#1e40af; font-weight:bold;">ERNIE 在线</div>
        </div>
        """, unsafe_allow_html=True)
        
    st.markdown("---")
    st.caption("Powered by Baidu Qianfan | Ver 9.5")

# --- 5. 主界面 ---

# 顶部 Banner 区域
st.markdown("""
<div style="padding: 20px; background: linear-gradient(90deg, #1e293b 0%, #334155 100%); border-radius: 12px; margin-bottom: 25px; color: white;">
    <h1 style="color: white; margin:0; font-size: 2.2rem;">文心·刑辩智脑</h1>
    <p style="opacity: 0.8; margin-top: 5px; font-size: 1rem;">基于大模型的全流程庭审推演与策略辅助系统</p>
</div>
""", unsafe_allow_html=True)

left_col, right_col = st.columns([4, 6], gap="large")

# === 左侧输入 ===
with left_col:
    st.subheader("📝 案情事实录入")
    # 使用 Container 增加白色背景卡片感
    with st.container():
        user_input = st.text_area(
            label="输入案情", 
            label_visibility="collapsed", 
            key="current_fact", 
            height=600, 
            placeholder="请在此输入案件详情，支持直接粘贴起诉书内容..."
        )
    
    st.markdown("<div style='height: 10px'></div>", unsafe_allow_html=True)
    
    c1, c2 = st.columns([2, 1])
    with c1:
        analyze_btn = st.button("🚀 启动深度分析", type="primary", use_container_width=True)
    with c2:
        st.button("🗑️ 清空", on_click=on_clear_click, type="secondary", use_container_width=True)

# === 右侧输出 ===
with right_col:
    st.subheader("📊 智能分析报告")
    
    # 辅助函数：随机抽取素材
    def get_few_shot_prompt():
        if not materials_content: return ""
        cases = materials_content.split('素材 No.')
        cases = [c for c in cases if len(c) > 50]
        if len(cases) > 0:
            selected_examples = random.sample(cases, min(2, len(cases)))
            return "请参考以下【优秀庭审辩论范例】的风格：\n" + "\n\n".join(selected_examples)
        return ""

    # A. 首次分析逻辑 (保持原有逻辑不变)
    if analyze_btn and user_input:
        st.session_state.debate_history = [] 
        st.session_state.messages = [] 
        with st.spinner('⚖️ 正在进行事件检测与争议焦点提取...'):
            try:
                learning_material = get_few_shot_prompt()

                prompt_text = f"""
                你是一名资深刑辩律师。请根据用户提供的案情进行全维度分析。
                
                {learning_material}
                
                ---
                【当前案情】：{user_input}
                
                【任务要求】：
                1. **事件检测**：梳理案情的时间线，提取关键节点。
                2. **争议焦点识别**：归纳本案控辩双方可能存在分歧的核心法律点（Controversy Focus）。
                3. **庭审推演**：模拟一场高水平的法庭辩论（至少3回合）。
                4. **量刑预测**：给出具体数值，严禁输出"待判决"。
                
                请严格按照以下JSON格式输出：
                {{
                    "charge": "罪名",
                    "term": "具体刑期预测（如：有期徒刑6个月至1年）",
                    "timeline": [
                        {{"time": "时间点/阶段", "event": "关键事件简述"}},
                        {{"time": "时间点/阶段", "event": "关键事件简述"}}
                    ],
                    "controversy": [
                        "争议焦点1（如：是否构成正当防卫）",
                        "争议焦点2（如：是否存在自首情节）"
                    ],
                    "analysis": "案情简析",
                    "strategy": "核心策略",
                    "law": "涉及法条（必须列出具体条文，如《中华人民共和国刑法》第XX条，并精简叙述每条的规则内容）",
                    "doc": "刑事辩护意见书全文",
                    "debate": [
                        {{"role": "公诉人", "content": "观点..."}},
                        {{"role": "辩护人", "content": "反驳..."}}
                    ]
                }}
                """
                resp = chat_comp.do(model="ERNIE-Speed-128K", messages=[{"role": "user", "content": prompt_text}], temperature=0.01)
                st.session_state.analysis_result = resp["body"]["result"]
                match = re.search(r'\{.*\}', st.session_state.analysis_result, re.DOTALL)
                if match:
                    try:
                        temp_data = json.loads(match.group(), strict=False)
                        st.session_state.debate_history = temp_data.get('debate', [])
                    except: pass
            except Exception as e:
                st.error(f"系统错误: {e}")

    # B. 结果渲染
    if st.session_state.analysis_result:
        result_text = st.session_state.analysis_result
        match = re.search(r'\{.*\}', result_text, re.DOTALL)
        data = {}
        if match:
            try:
                # 尝试标准解析
                clean_str = match.group().replace('\n', '\\n').replace('\t', '')
                data = json.loads(clean_str, strict=False)
            except: 
                try: 
                    # 尝试宽松解析
                    data = json.loads(match.group(), strict=False)
                except: 
                    try:
                        # 🔥 终极解析：使用 ast 处理单引号问题
                        data = ast.literal_eval(match.group())
                    except: pass

        if data:
            # 渲染卡片：使用新的 CSS 类
            def info_card(title, content, border_color="#1e293b"):
                return f"""
                <div class="result-card" style="border-left: 5px solid {border_color};">
                    <div class="card-title">{title}</div>
                    <div class="card-value">{content}</div>
                </div>
                """

            m1, m2 = st.columns(2)
            with m1: st.markdown(info_card("⚖️ 预测罪名", data.get('charge', '未知'), "#1e293b"), unsafe_allow_html=True)
            with m2: st.markdown(info_card("📅 建议量刑", data.get('term', '正在评估...'), "#b49162"), unsafe_allow_html=True)
            st.write("") 

            # 🔥 Tab 布局
            tab1, tab2, tab3, tab4, tab5 = st.tabs(["⚔️ 庭审推演", "🧠 深度分析", "📜 辩护词生成", "📖 法律依据", "📂 类案推荐"])
            
            # === Tab 1: 滚动版辩论战场 ===
            with tab1:
                st.markdown("##### 🏛️ 模拟法庭辩论实录")
                st.caption("AI 仿真生成 · 控辩双方实时对抗")
                
                with st.container(height=500, border=True):
                    for turn in st.session_state.debate_history:
                        role = turn.get('role', '')
                        content = turn.get('content', '')
                        if "公诉" in role:
                            with st.chat_message("user", avatar="🏛️"): 
                                st.markdown(f"**【{role}】**")
                                st.write(content)
                        else:
                            with st.chat_message("assistant", avatar="🛡️"):
                                st.markdown(f"**【{role}】**")
                                st.write(content)

                st.write("") 
                
                # 🔥🔥🔥 续写按钮 🔥🔥🔥
                if st.button("🔥 继续激烈交锋 (Extend Debate)", type="secondary", use_container_width=True):
                    with st.spinner("⏳ 控辩双方正在组织新一轮观点..."):
                        # ... (续写逻辑保持不变) ...
                        learning_material = get_few_shot_prompt()
                        history_context = json.dumps(st.session_state.debate_history, ensure_ascii=False)
                        
                        extend_prompt = f"""
                        {learning_material}
                        ---
                        【当前案情】：{user_input}
                        【目前的庭审辩论记录】：{history_context}
                        【任务】：请继续生成下一轮交锋（Round 2）。
                        1. 公诉人针对辩护人刚才的观点进行有力反驳。
                        2. 辩护人再次进行防御。
                        【注意】：必须保持上述范例中的专业、犀利风格。
                        请严格只输出一个JSON列表（不要Markdown）：
                        [
                            {{"role": "公诉人", "content": "..."}},
                            {{"role": "辩护人", "content": "..."}}
                        ]
                        """
                        try:
                            resp_next = chat_comp.do(model="ERNIE-Speed-128K", messages=[{"role": "user", "content": extend_prompt}], temperature=0.1)
                            new_content = resp_next["body"]["result"]
                            
                            match_list = re.search(r'\[.*\]', new_content, re.DOTALL)
                            if match_list:
                                new_turns = []
                                list_str = match_list.group()
                                try:
                                    new_turns = json.loads(list_str, strict=False)
                                except:
                                    try:
                                        clean_str = list_str.replace('\n', ' ').replace('\t', '')
                                        new_turns = json.loads(clean_str, strict=False)
                                    except:
                                        try:
                                            new_turns = ast.literal_eval(list_str)
                                        except: pass
                                
                                if new_turns:
                                    st.session_state.debate_history.extend(new_turns)
                                    st.rerun()
                                else:
                                    st.warning("续写格式异常，请重试")
                        except Exception as e:
                            st.warning(f"续写遇到一点小插曲 ({e})")

            # === Tab 2: 深度分析 ===
            with tab2:
                col_focus, col_timeline = st.columns([1, 1])
                
                with col_focus:
                    st.markdown("##### 🔥 核心争议焦点")
                    controversies = data.get('controversy', [])
                    if controversies:
                        for i, point in enumerate(controversies):
                            st.info(f"**焦点 {i+1}：** {point}")
                    else:
                        st.info(data.get('analysis', ''))
                    
                    st.markdown("##### 🛡️ 核心辩护策略")
                    st.success(data.get('strategy', ''))

                with col_timeline:
                    st.markdown("##### ⏱️ 案情还原时间轴")
                    timelines = data.get('timeline', [])
                    if timelines:
                        # 🔥 修复点：去掉 f-string 内部的缩进，防止被识别为代码块
                        timeline_html = '<div class="timeline-container">'
                        for item in timelines:
                            timeline_html += f"""
                            <div class="timeline-item">
                                <div class="timeline-time">{item.get('time', '')}</div>
                                <div class="timeline-content">{item.get('event', '')}</div>
                            </div>"""
                        timeline_html += '</div>'
                        st.markdown(timeline_html, unsafe_allow_html=True)
                    else:
                        st.text("暂无时间轴数据")

            with tab3:
                st.markdown("##### 📜 《刑事辩护意见书》草案")
                st.caption("AI 自动生成的初稿，仅供律师参考")
                st.text_area("可复制文本", value=data.get('doc', '正在生成...'), height=500)
                
            with tab4:
                st.markdown("##### 📖 法律依据详解")
                st.warning(data.get('law', ''))

            with tab5:
                st.markdown("##### 📂 相似判例推荐")
                if not df.empty:
                    predicted_charge = data.get('charge', '').replace("罪", "")
                    similar_cases = df[df['real_charge'].str.contains(predicted_charge, na=False)]
                    if not similar_cases.empty:
                        cnt = 0
                        for idx, row in similar_cases.iterrows():
                            if str(row['fact']).strip() != user_input.strip():
                                st.markdown(f"""
                                <div style="padding:15px; border:1px solid #e2e8f0; border-radius:8px; margin-bottom:10px; background:white;">
                                    <div style="font-weight:bold; color:#1e293b;">案例 #{idx} - {row['real_charge']}</div>
                                    <div style="color:#64748b; font-size:0.9em; margin:5px 0;">判决结果：<span style="color:#b49162; font-weight:bold;">{row['real_term']} 个月</span></div>
                                    <div style="color:#334155; font-size:0.9em;">{str(row['fact'])[:80]}...</div>
                                </div>
                                """, unsafe_allow_html=True)
                                cnt += 1
                            if cnt >= 2: break
                    else:
                        st.text("暂无高度匹配案例")
                else:
                    st.text("数据库未连接")

            # === 3. 底部交互问答区 ===
            st.divider()
            st.subheader("💬 AI 案情助手")
            
            for message in st.session_state.messages:
                with st.chat_message(message["role"]):
                    st.markdown(message["content"])

            if prompt := st.chat_input("针对本案，您还有什么想问的？"):
                st.session_state.messages.append({"role": "user", "content": prompt})
                with st.chat_message("user"):
                    st.markdown(prompt)

                with st.chat_message("assistant"):
                    with st.spinner("AI 正在检索法律库..."):
                        context_prompt = f"""
                        你是一名刑辩律师助手。
                        当前正在分析的案情：{user_input}
                        此前的分析结论：{data.get('analysis')}
                        用户现在的提问：{prompt}
                        请结合案情和法律规定进行简明扼要的回答。
                        """
                        try:
                            resp_chat = chat_comp.do(model="ERNIE-Speed-128K", messages=[{"role": "user", "content": context_prompt}], temperature=0.01)
                            ai_reply = resp_chat["body"]["result"]
                            st.markdown(ai_reply)
                            st.session_state.messages.append({"role": "assistant", "content": ai_reply})
                        except Exception as e:
                            st.error(f"对话出错: {e}")

    elif not analyze_btn:
        # 空状态提示美化
        st.markdown("""
        <div style="text-align: center; padding: 60px; color: #94a3b8; background: #ffffff; border-radius: 12px; border: 2px dashed #e2e8f0;">
            <div style="font-size: 40px; margin-bottom: 20px;">⚖️</div>
            <div style="font-size: 18px; font-weight: bold; color: #475569;">等待案情输入</div>
            <div style="font-size: 14px; margin-top: 5px;">请在左侧输入案情详情，并点击“启动深度分析”</div>
        </div>
        """, unsafe_allow_html=True)